package easy.trace.dubbo;


import com.alibaba.dubbo.common.extension.Activate;
import com.alibaba.dubbo.remoting.TimeoutException;
import com.alibaba.dubbo.rpc.Filter;
import com.alibaba.dubbo.rpc.Invocation;
import com.alibaba.dubbo.rpc.Invoker;
import com.alibaba.dubbo.rpc.Result;
import com.alibaba.dubbo.rpc.RpcContext;
import com.alibaba.dubbo.rpc.RpcException;
import com.alibaba.dubbo.rpc.RpcInvocation;
import easy.trace.BinaryAnnotation;
import easy.trace.Endpoint;
import easy.trace.Span;
import easy.trace.agent.Tracer;
import easy.trace.agent.support.TracerUtils;
import java.io.IOException;
import java.util.Properties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.support.PropertiesLoaderUtils;

@Activate(
        group = {"provider", "consumer"}
)
public class HydraFilter implements Filter {
    private static Logger logger = LoggerFactory.getLogger(HydraFilter.class);
    private static Tracer tracer = null;
    private static boolean isTraceStart;

    public HydraFilter() {
    }

    private boolean isTestEnv() {
        try {
            ClassPathResource resource = new ClassPathResource("/config.properties");
            Properties props = PropertiesLoaderUtils.loadProperties(resource);
            String env = (String)props.get("env");
            if (env != null) {
                return "dev".equals(env) || "prod".equals(env) || "stage".equals(env) || "integration".equals(env);
            } else {
                return false;
            }
        } catch (IOException var4) {
            logger.error(var4.getMessage());
            return false;
        }
    }

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        if (tracer != null) {
            Tracer.getTracer();
            if (!Tracer.isTraceStartedFail()) {
                String serviceId = tracer.getServiceId(RpcContext.getContext().getUrl().getServiceInterface());
                long start = System.currentTimeMillis();
                RpcContext context = RpcContext.getContext();
                boolean isConsumerSide = context.isConsumerSide();
                Span span = null;
                Endpoint endpoint = null;
                boolean var20 = false;

                Result var26;
                try {
                    var20 = true;
                    endpoint = tracer.newEndPoint();
                    endpoint.setIp(context.getLocalAddressString());
                    endpoint.setPort(context.getLocalPort());
                    Long traceId;
                    if (context.isConsumerSide()) {
                        Span span1 = tracer.getParentSpan();
                        if (span1 == null) {
                            traceId = null;
                            span = tracer.newSpan(serviceId, context.getMethodName(), endpoint, traceId);
                        } else {
                            span = tracer.genSpan(span1.getTraceId(), span1.getId(), tracer.genSpanId(), context.getMethodName(), span1.isSample(), serviceId);
                        }
                    } else if (context.isProviderSide()) {
                        traceId = TracerUtils.getAttachmentLong(invocation.getAttachment("traceId"));
                        traceId = TracerUtils.getAttachmentLong(invocation.getAttachment("parentId"));
                        Long spanId = TracerUtils.getAttachmentLong(invocation.getAttachment("spanId"));
                        boolean isSample = traceId != null;
                        span = tracer.genSpan(traceId, traceId, spanId, context.getMethodName(), isSample, serviceId);
                    }

                    this.invokerBefore(invocation, span, endpoint, start);
                    RpcInvocation invocation1 = (RpcInvocation)invocation;
                    this.setAttachment(span, invocation1);
                    Result result = invoker.invoke(invocation);
                    if (result.getException() != null) {
                        this.catchException(result.getException(), endpoint);
                    }

                    var26 = result;
                    var20 = false;
                } catch (RpcException var21) {
                    if (var21.getCause() != null && var21.getCause() instanceof TimeoutException) {
                        this.catchTimeoutException(var21, endpoint);
                    } else {
                        this.catchException(var21, endpoint);
                    }

                    throw var21;
                } finally {
                    if (var20) {
                        if (span != null) {
                            long end = System.currentTimeMillis();
                            this.invokerAfter(invocation, endpoint, span, end, isConsumerSide);
                        }

                    }
                }

                if (span != null) {
                    long end = System.currentTimeMillis();
                    this.invokerAfter(invocation, endpoint, span, end, isConsumerSide);
                }

                return var26;
            }
        }

        tracer = Tracer.getTracer();
        return invoker.invoke(invocation);
    }

    private void catchTimeoutException(RpcException e, Endpoint endpoint) {
        BinaryAnnotation exAnnotation = new BinaryAnnotation();
        exAnnotation.setKey("dubbo.exception");
        exAnnotation.setValue(e.getMessage());
        exAnnotation.setType("exTimeout");
        exAnnotation.setHost(endpoint);
        tracer.addBinaryAnntation(exAnnotation);
    }

    private void catchException(Throwable e, Endpoint endpoint) {
        BinaryAnnotation exAnnotation = new BinaryAnnotation();
        exAnnotation.setKey("dubbo.exception");
        exAnnotation.setValue(e.getMessage());
        exAnnotation.setType("ex");
        exAnnotation.setHost(endpoint);
        tracer.addBinaryAnntation(exAnnotation);
    }

    private void setAttachment(Span span, RpcInvocation invocation) {
        if (span.isSample()) {
            invocation.setAttachment("parentId", span.getParentId() != null ? String.valueOf(span.getParentId()) : null);
            invocation.setAttachment("spanId", span.getId() != null ? String.valueOf(span.getId()) : null);
            invocation.setAttachment("traceId", span.getTraceId() != null ? String.valueOf(span.getTraceId()) : null);
        }

    }

    private void invokerAfter(Invocation invocation, Endpoint endpoint, Span span, long end, boolean isConsumerSide) {
        if (isConsumerSide && span.isSample()) {
            tracer.clientReceiveRecord(span, endpoint, end);
        } else {
            tracer.serverSendRecord(span, endpoint, end);
            tracer.removeParentSpan();
        }

    }

    private void invokerBefore(Invocation invocation, Span span, Endpoint endpoint, long start) {
        RpcContext context = RpcContext.getContext();
        if (context.isConsumerSide() && span.isSample()) {
            tracer.clientSendRecord(span, endpoint, start);
        } else if (context.isProviderSide()) {
            tracer.serverReceiveRecord(span, endpoint, start);
            tracer.setParentSpan(span);
        }

    }

    public static Tracer getTracer() {
        return tracer;
    }

    public static void setTracer(Tracer tracer) {
        tracer = tracer;
    }
}
